{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bigbird import modeling\n",
    "from bigbird import utils\n",
    "import tensorflow as tf\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "bert_config = {\n",
    "  # transformer basic configs\n",
    "  \"attention_probs_dropout_prob\": 0.1,\n",
    "  \"hidden_act\": 'gelu',\n",
    "  \"hidden_dropout_prob\": 0.1,\n",
    "  \"hidden_size\": 768,\n",
    "  \"initializer_range\": 0.02,\n",
    "  \"intermediate_size\": 3072,\n",
    "  \"max_position_embeddings\": 4096,\n",
    "  \"max_encoder_length\": 1024,\n",
    "  \"num_attention_heads\": 12,\n",
    "  \"num_hidden_layers\": 12,\n",
    "  \"type_vocab_size\": 2,\n",
    "  \"scope\": 'bert',\n",
    "  \"use_bias\": True,\n",
    "  \"rescale_embedding\": False,\n",
    "  \"vocab_model_file\": None,\n",
    "  # sparse mask configs\n",
    "  \"attention_type\": \"block_sparse\",\n",
    "  \"norm_type\": 'postnorm',\n",
    "  \"block_size\": 16,\n",
    "  \"num_rand_blocks\": 3,\n",
    "  \"vocab_size\": 32000,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = modeling.BertModel(bert_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = tf.placeholder(tf.int32, [None, None])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /Users/huseinzolkepli/Documents/Malaya/pretrained-model/bigbird/bigbird/modeling.py:124: where (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.where in 2.0, which has the same broadcast rule as np.where\n"
     ]
    }
   ],
   "source": [
    "sequence_output, pooled_output = model(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<tf.Tensor 'bert/encoder/layer_11/output/LayerNorm/batchnorm/add_1:0' shape=(?, 1024, 768) dtype=float32>,\n",
       " <tf.Tensor 'bert/pooler/dense/Tanh:0' shape=(?, 768) dtype=float32>)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sequence_output, pooled_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "sess = tf.InteractiveSession()\n",
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import collections\n",
    "import re\n",
    "\n",
    "def get_assignment_map_from_checkpoint(tvars, init_checkpoint):\n",
    "    \"\"\"Compute the union of the current variables and checkpoint variables.\"\"\"\n",
    "    assignment_map = {}\n",
    "    initialized_variable_names = {}\n",
    "\n",
    "    name_to_variable = collections.OrderedDict()\n",
    "    for var in tvars:\n",
    "        name = var.name\n",
    "        m = re.match('^(.*):\\\\d+$', name)\n",
    "        if m is not None:\n",
    "            name = m.group(1)\n",
    "        name_to_variable[name] = var\n",
    "\n",
    "    init_vars = tf.train.list_variables(init_checkpoint)\n",
    "\n",
    "    assignment_map = collections.OrderedDict()\n",
    "    for x in init_vars:\n",
    "        (name, var) = (x[0], x[1])\n",
    "        name_r = name.replace('bert/embeddings/LayerNorm', 'bert/encoder/LayerNorm')\n",
    "        if name_r not in name_to_variable:\n",
    "            continue\n",
    "        if 'embeddings/position_embeddings' in name_r:\n",
    "            continue\n",
    "        assignment_map[name] = name_to_variable[name_r]\n",
    "        initialized_variable_names[name_r] = 1\n",
    "        initialized_variable_names[name_r + ':0'] = 1\n",
    "\n",
    "    return (assignment_map, initialized_variable_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !wget https://f000.backblazeb2.com/file/malaya-model/bert-bahasa/bert-base-2020-10-08.tar.gz\n",
    "# !tar -zxf bert-base-2020-10-08.tar.gz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "tvars = tf.trainable_variables()\n",
    "checkpoint = 'bert-base/model.ckpt-1000000'\n",
    "assignment_map, initialized_variable_names = get_assignment_map_from_checkpoint(tvars, \n",
    "                                                                                checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Restoring parameters from bert-base/model.ckpt-1000000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Restoring parameters from bert-base/model.ckpt-1000000\n"
     ]
    }
   ],
   "source": [
    "saver = tf.train.Saver(var_list = assignment_map)\n",
    "saver.restore(sess, checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 6.43 s, sys: 1.16 s, total: 7.6 s\n",
      "Wall time: 1.28 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "o = sess.run([sequence_output, pooled_output], feed_dict = {X: [[1] * 1024]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
